package org.nicta.datasetLoader.MovieLens;

import java.io.BufferedReader;
import java.io.FileReader;
import java.util.HashMap;
import java.util.HashSet;

public class MovieLensDataLoader extends DataLoader {
		
    private String ratingsFileLocation = null;
    private String usersFileLocation = null;
    private String moviesFileLocation = null;

	private HashMap<Integer[], Double> ratings;
	private HashMap<Integer, double[]> userFeatures;
    private HashMap<Integer, double[]> movieFeatures;

	private MovieLensInfo info;
	private boolean isDataLoaded = false;
	
	public MovieLensDataLoader(String datasetType, String ratingsFileLocation,
			String usersFileLocation,String moviesFileLocation){
		
		this.ratingsFileLocation = ratingsFileLocation;
		this.usersFileLocation = usersFileLocation;
		this.moviesFileLocation = moviesFileLocation;
		
		info = MovieLensHelper.getInfo(datasetType);
	}
		
	private void loadRatings() throws Exception{
		if(ratingsFileLocation == null){
			if(getDebugMode()){
				System.out.println("Error: Ratings File Location Not " +
									"Specified");
			}
			return;
		}		
		BufferedReader ratingReader = 
				new BufferedReader(new FileReader(ratingsFileLocation));
		
		if(info.dataSetType.equals(MovieLensHelper.small)){
			ratings = MovieLensHelper.load100KRatings(ratingReader,info);
		}
		else if(info.dataSetType.equals(MovieLensHelper.medium))
		{
			ratings = MovieLensHelper.load1MRatings(ratingReader,info);

		}
		else if(info.dataSetType.equals(MovieLensHelper.large)){
			ratings = MovieLensHelper.load10MRatings(ratingReader,info);	
		}
		ratingReader.close();


	}
	
	private void loadUserFeatures() throws Exception{
		if(usersFileLocation == null){
			if(getDebugMode()){
				System.out.println("Error: Users File Location Not " +
									"Specified");
			}
			return;
		}
       
		BufferedReader userFeaturesReader = new BufferedReader(
											new FileReader(usersFileLocation));
        
		if(info.dataSetType.equals(MovieLensHelper.small)){
			userFeatures = 
					MovieLensHelper.load100KUserFeatures(userFeaturesReader,info);
		}
		else if(info.dataSetType.equals(MovieLensHelper.medium))
		{
			userFeatures = 
					MovieLensHelper.load1MUserFeatures(userFeaturesReader,info);
		}
		else if(info.dataSetType.equals(MovieLensHelper.large)){
			userFeatures = 
					MovieLensHelper.load10MUserFeatures(userFeaturesReader,info);	
		}
        userFeaturesReader.close();
	}
	
	
	private void loadMovieFeatures() throws Exception{
		if(moviesFileLocation == null){
			if(getDebugMode()){
				System.out.println("Error: Users File Location Not " +
									"Specified");
			}
			return;
		}
       
		BufferedReader movieFeaturesReader = new BufferedReader(
				new FileReader(moviesFileLocation));
		
		if(info.dataSetType.equals(MovieLensHelper.small)){
		movieFeatures = 
				MovieLensHelper.load100KMovieFeatures(movieFeaturesReader,info);
		}
		else if(info.dataSetType.equals(MovieLensHelper.medium))
		{
		movieFeatures = 
				MovieLensHelper.load1MMovieFeatures(movieFeaturesReader,info);
		}
		else if(info.dataSetType.equals(MovieLensHelper.large)){
		movieFeatures = 
				MovieLensHelper.load10MMovieFeatures(movieFeaturesReader,info);	
		}
		movieFeaturesReader.close();
	}
	
	@Override
	public void loadData() throws Exception {
		// TODO Auto-generated method stub
		loadRatings();
		loadUserFeatures();
		loadMovieFeatures();
		isDataLoaded = true;
	}
	
	public HashMap<Integer, double[]> getMoviesFeatures(){
		return movieFeatures;
	}
	
	public HashMap<Integer, double[]> getUsersFeatures(){
		return userFeatures;
	}
	
	public HashMap<Integer[], Double> getRatings(){
		return ratings;
	}
	
	public HashMap<Integer, HashMap<Integer, Double>> getmovieUserRatings(){
		HashMap<Integer, HashMap<Integer, Double>> movieUserRatings =
				new HashMap<Integer, HashMap<Integer, Double>>();
		if(!isDataLoaded)
			return null;
		Object[] ratingsKey  = ratings.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			Double rating =  ratings.get(userItemPair);
		    if (!movieUserRatings.containsKey(movieId)) {
		    	movieUserRatings.put(movieId, new HashMap<Integer, Double>());
		    }
            movieUserRatings.get(movieId).put(userId, rating);
		}
		return movieUserRatings;
	}
	

	public HashMap<Integer, HashSet<Integer>> getUserRatedMoviesMap(){
		HashMap<Integer, HashSet<Integer>> userRatedMovies = 
				new HashMap<Integer, HashSet<Integer>>();
		
		if(!isDataLoaded)
			return null;
		
		Object[] ratingsKey  = ratings.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			if(!userRatedMovies.containsKey(userId)){
				HashSet<Integer> movies = new HashSet<Integer>(movieId);
				userRatedMovies.put(userId, movies);
			}
			userRatedMovies.get(userId).add(movieId);
		}
		return userRatedMovies;
		
	}
	
	public HashMap<Integer, HashSet<Integer>> getMoviesUsersMap(){
		HashMap<Integer, HashSet<Integer>> moviesUserMap = 
				new HashMap<Integer, HashSet<Integer>>();
		if(!isDataLoaded)
			return null;
		
		for(Integer[] userItemPair: ratings.keySet()){
			Integer userId = userItemPair[0];
			Integer movieId = userItemPair[1];
			if(!moviesUserMap.containsKey(movieId)){
				HashSet<Integer> users = new HashSet<Integer>(userId);
				moviesUserMap.put(movieId, users);
			}
			else{
				moviesUserMap.get(movieId).add(userId);
			}
		}
		return moviesUserMap;
	}
	
	public Integer getUserFeaturesSize(){
		return info.userFeaturesCount;
	}
	
	public Integer getMovieFeaturesSize(){
		return info.movieFeaturesCount;
	}
	
	public Integer getUsersCount(){
		return info.usersCount;
	}
	
	public Integer getMoviesCount(){
		return info.moviesCount;
	}
 
}
